USE 113 - Rework model registration and loading#13
Merged
Conversation
Why these changes are being introduced: Previously, a model's URI was only codified in the registry.py file, and not directly obvious from a class itself aside from docstrings. How this addresses that need: Reworks the base class to require a MODEL_URI attribute. This attribute is used to generate a model_uri:class dictionary that is used for registration of models. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-113
ghukill
commented
Oct 24, 2025
Comment on lines
+13
to
+24
| MODEL_URI: str # Type hint to document the requirement | ||
|
|
||
| def __init_subclass__(cls, **kwargs: dict) -> None: # noqa: D105 | ||
| super().__init_subclass__(**kwargs) | ||
|
|
||
| # require class level MODEL_URI to be set | ||
| if not hasattr(cls, "MODEL_URI"): | ||
| msg = f"{cls.__name__} must define 'MODEL_URI' class attribute" | ||
| raise TypeError(msg) | ||
| if not isinstance(cls.MODEL_URI, str): | ||
| msg = f"{cls.__name__} must override 'MODEL_URI' with a valid string" | ||
| raise TypeError(msg) |
Collaborator
Author
There was a problem hiding this comment.
This establishes MODEL_URI as a required class-level attribute for all child classes. There are different ways to achieve this, but this felt like a decent, explicit pattern here.
ghukill
commented
Oct 24, 2025
| HuggingFace URI: opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte | ||
| """ | ||
|
|
||
| MODEL_URI = "opensearch-project/opensearch-neural-sparse-encoding-doc-v3-gte" |
Collaborator
Author
There was a problem hiding this comment.
Now, each model (and we may only have a couple ever), will define their HuggingFace URI in the class itself.
ghukill
commented
Oct 24, 2025
Comment on lines
+16
to
+29
| def model_required(f: Callable) -> Callable: | ||
| """Decorator for commands that require a specific model.""" | ||
|
|
||
| @click.option( | ||
| "--model-uri", | ||
| envvar="TE_MODEL_URI", | ||
| required=True, | ||
| help="HuggingFace model URI (e.g., 'org/model-name')", | ||
| ) | ||
| @functools.wraps(f) | ||
| def wrapper(*args: list, **kwargs: dict) -> Callable: | ||
| return f(*args, **kwargs) | ||
|
|
||
| return wrapper |
Collaborator
Author
There was a problem hiding this comment.
DRY's up the --model-uri CLI arg, and gives it an env var default.
ehanson8
approved these changes
Oct 27, 2025
ehanson8
left a comment
ehanson8
left a comment
There was a problem hiding this comment.
Smart update and works as expected!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose and background context
Before diving fully into the our first model
OSNeuralSparseDocV3GTEbeing capable of downloading the model locally, a bit of reworking for how embedding model classes are discovered and loaded.Why these changes are being introduced:
Previously, a model's URI was only codified in the
embeddings/registry.pyfile, and not directly obvious from a class itself aside from docstrings.Additionally, the CLI argument
--model-uriwas duplicated across multiple CLI commands, where we can DRY that up and provide an environment variable default.How this addresses that need:
Reworks the base class to require a
MODEL_URIattribute. This attribute is used to generate amodel_uri:classdictionary that is used for registration of models.How can a reviewer manually see the effects of these changes?
1- Create a
.envfile with the following:2- Run the CLI command
download-model. While an error is raised because not yet implemented, note that--model-urior--outputwere not required due to env vars getting set:NOTE: the use of
uv run --env-file .env .... As noted in this uv github issue, there is lots of interest in making this default behavior... but it has not yet landed inuv. For the time being, it sounds like we may need to somewhat explicitly request thatuvread the.envfile when running locally.In a deployed context, the env vars will be set and this is not required. To that end, we could also set the env vars in the terminal and just run
embeddings ...; dealer's choice.Includes new or updated dependencies?
YES
Changes expectations for external applications?
NO
What are the relevant tickets?
Code review